package com.huawei.admins.platform.gateway.core.component.limiter;

import com.huawei.admins.platform.gateway.core.component.filter.factory.RequestRateLimiterGatewayFilterFactory.CompositeConfig;
import com.huawei.admins.platform.gateway.core.component.filter.factory.RequestRateLimiterGatewayFilterFactory.Config;
import com.huawei.admins.platform.gateway.core.config.GatewayProperties;
import org.apache.commons.lang3.tuple.Pair;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class ServiceRateLimiter extends AbstractRateLimiter {

    public static final String CONFIGURATION_PROPERTY_NAME = "local-rate-limiter";

    private final Map<String, ServiceRateLimiterTool> serviceRateLimiterMap = new ConcurrentHashMap<>();

    private final RedisRateLimiter redisRateLimiter;

    public ServiceRateLimiter(RedisRateLimiter redisRateLimiter) {
        this.redisRateLimiter = redisRateLimiter;
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id, CompositeConfig config) {

        String limiterKey = routeId + ":" + id;
        Config routeConfig = config.getLocalConfig();

        ServiceRateLimiterTool serviceRateLimiter = serviceRateLimiterMap.computeIfAbsent(routeId, key -> {
            int replenishRate = routeConfig.getReplenishRate();
            int burstCapacity = routeConfig.getBurstCapacity();
            int requestedTokens = routeConfig.getRequestedTokens();
            return new ServiceRateLimiterTool(replenishRate, burstCapacity, requestedTokens);
        });

        Pair<Boolean, Long> result = serviceRateLimiter.tryAcquire(limiterKey);
        if (result.getLeft()) {
            return Mono.just(new Response(true, getHeaders(routeConfig, result.getRight())));
        }
        return redisRateLimiter.isAllowed(routeId, id, config);
    }
}
